#!/usr/bin/env python3
"""
下载所有TTS模型到指定目录
"""

import os
import sys
import logging

# 设置环境变量
os.environ['XDG_CACHE_HOME'] = os.path.expanduser('~/.cache')
os.environ['XDG_DATA_HOME'] = os.path.expanduser('~/.local/share')
os.environ['TTS_CACHE_DIR'] = os.path.expanduser('~/.cache/tts')
os.environ['TORCH_HOME'] = os.path.expanduser('~/.cache/torch')
os.environ['HF_HOME'] = os.path.expanduser('~/.cache/huggingface')

# 创建目录
os.makedirs(os.path.expanduser('~/.cache/tts'), exist_ok=True)
os.makedirs(os.path.expanduser('~/.cache/torch'), exist_ok=True)
os.makedirs(os.path.expanduser('~/.cache/huggingface'), exist_ok=True)

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

try:
    from TTS.api import TTS
    print("TTS库导入成功")
except ImportError as e:
    print(f"TTS库导入失败: {e}")
    sys.exit(1)

# 定义要下载的模型列表
models_to_download = [
    # 中文模型
    ("tts_models/zh-CN/baker/tacotron2-DDC-GST", "Baker - 中文专用模型"),
    
    # 多语言模型
    ("tts_models/multilingual/multi-dataset/bark", "Bark - 最自然的多语言TTS"),
    ("tts_models/multilingual/multi-dataset/your_tts", "YourTTS - 高质量多语言模型"),
    ("tts_models/multilingual/multi-dataset/xtts_v2", "XTTS v2 - 最先进的克隆语音"),
    
    # 英文模型
    ("tts_models/en/ljspeech/vits", "VITS - 高质量英文模型"),
    ("tts_models/en/ljspeech/tacotron2-DDC", "Tacotron2 - 稳定英文模型"),
]

def download_model(model_name, description):
    """下载单个模型"""
    try:
        logger.info(f"开始下载: {description}")
        logger.info(f"模型名称: {model_name}")
        
        # 创建TTS实例，这会自动下载模型
        tts = TTS(model_name=model_name, progress_bar=True)
        
        logger.info(f"✅ 下载完成: {description}")
        return True
        
    except Exception as e:
        logger.error(f"❌ 下载失败: {description}")
        logger.error(f"错误信息: {e}")
        return False

def main():
    """主函数"""
    print("="*60)
    print("🚀 开始下载TTS模型到 ~/.cache/tts/")
    print("="*60)
    
    success_count = 0
    total_count = len(models_to_download)
    
    for model_name, description in models_to_download:
        print(f"\n📦 正在处理: {description}")
        print(f"🔗 模型路径: {model_name}")
        
        if download_model(model_name, description):
            success_count += 1
        
        print("-" * 50)
    
    print(f"\n✨ 下载完成统计:")
    print(f"✅ 成功: {success_count}/{total_count}")
    print(f"❌ 失败: {total_count - success_count}/{total_count}")
    
    if success_count == total_count:
        print("\n🎉 所有模型下载成功！")
    else:
        print(f"\n⚠️  有 {total_count - success_count} 个模型下载失败")
    
    # 显示下载位置
    cache_dir = os.path.expanduser('~/.cache/tts')
    print(f"\n📁 模型保存位置: {cache_dir}")
    
    # 列出已下载的模型
    try:
        if os.path.exists(cache_dir):
            print("\n📋 已下载的模型:")
            for item in os.listdir(cache_dir):
                if os.path.isdir(os.path.join(cache_dir, item)):
                    print(f"  📂 {item}")
    except Exception as e:
        print(f"列出模型时出错: {e}")

if __name__ == "__main__":
    main()